import requests
class ModelSetting(object):
    def __init__(self,apiKey,secretKey):
        self.SetApiKey(apiKey)
        self.SetSecretKey(secretKey)
    def SetApiKey(self,apiKey):
        self.apiKey = apiKey
    def SetSecretKey(self, secretKey):
        self.secretKey = secretKey
    def GetAccessToken(self):
        TokenUrl = f"https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id={self.apiKey}&client_secret={self.secretKey}"
        response = requests.get(TokenUrl)
        AccessToken = response.json()["access_token"]
        return AccessToken
    def GetChatApiUrl(self,modelName):
        AccessToken = self.GetAccessToken()
        if modelName == 'ERNIE-Bot-turbo':
            ChatApiUrl = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/eb-instant?access_token={AccessToken}"
        elif modelName == 'BLOOMZ-7B':
            ChatApiUrl = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/bloomz_7b1?access_token={AccessToken}"
        elif modelName == 'Embedding-V1':
            ChatApiUrl = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/embeddings/embedding-v1?access_token={AccessToken}"
        else:
            ChatApiUrl = f"https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token={AccessToken}"
        return ChatApiUrl